"""
特征工程脚本
功能：数据读取、预处理、特征转换和数据输出
"""

import pandas as pd
import config


def inputdata(path, nrows=None):
    """
    读取CSV数据文件
    
    Args:
        path (str): CSV文件路径
        nrows (int, optional): 读取的行数限制
    
    Returns:
        pd.DataFrame: 读取的数据框
    """
    data = pd.read_csv(path, header=0, sep=",", encoding="utf-8", nrows=nrows)
    return data


def outputdata(path, data, is_index=False):
    """
    输出数据到CSV文件
    
    Args:
        path (str): 输出文件路径
        data (pd.DataFrame): 要输出的数据框
        is_index (bool): 是否包含索引列
    """
    data.to_csv(path, index=is_index, header=True, sep=",", mode="w", encoding="utf-8")


def transcolname(df, column_mapping):
    """
    转换列名为英文
    
    Args:
        df (pd.DataFrame): 输入数据框
        column_mapping (dict): 列名映射字典
    
    Returns:
        pd.DataFrame: 列名转换后的数据框
    """
    df.rename(columns=column_mapping, inplace=True)
    return df


def trans_datetime(df):
    """
    日期时间特征转换
    将日期字符串转换为数值特征，并提取年月日信息
    
    Args:
        df (pd.DataFrame): 包含Date列的数据框
    
    Returns:
        pd.DataFrame: 添加了时间特征的数据框
    """
    ret_df = pd.DataFrame()
    dt = df["Date"]
    
    # 提取年月日特征
    ret_df["year"] = dt.transform(lambda x: int(x.split("-")[0]))
    ret_df["month"] = dt.transform(lambda x: int(x.split("-")[1]))
    ret_df["day"] = dt.transform(lambda x: int(x.split("-")[2][:2]))
    
    # 合并时间特征到原数据框
    df = pd.concat([df, ret_df], axis=1)
    
    # 将日期转换为序号（按时间顺序排列的序号）
    unique_dates = pd.Series(df["Date"].unique()).sort_values().reset_index(drop=True)
    date_mapping = {date: rank + 1 for rank, date in enumerate(unique_dates)}
    df["Date"] = df["Date"].map(date_mapping)
    
    # 备选方案：将日期转换为距离最早日期的天数
    # df["Date"] = pd.to_datetime(df["Date"], format="%Y-%m-%d")
    # minTime = df["Date"].min()
    # df["Date"] = ((df["Date"] - minTime) / pd.Timedelta(days=1)).astype(int)
    
    return df


def processing_feature():
    """
    主要特征处理函数
    读取训练数据，进行列名转换、特征工程等预处理
    
    Returns:
        pd.DataFrame: 处理后的特征数据
    """
    # 先读取少量数据获取总行数，用于计算采样比例
    temp_data = pd.read_csv("./data/train.csv", header=0, sep=",", encoding="utf-8")
    total_rows = len(temp_data)
    
    # 根据配置文件中的因子计算要读取的行数
    rows_to_read = int(total_rows * config.DATASET_FACTOR)
    
    # 读取指定数量的训练数据
    data = inputdata("./data/train.csv", nrows=rows_to_read)
    
    # 定义中文列名到英文列名的映射
    column_mapping = {
        "股票代码": "StockCode",      # 股票代码
        "日期": "Date",              # 日期
        "开盘": "Open",              # 开盘价
        "收盘": "Close",             # 收盘价
        "最高": "High",              # 最高价
        "最低": "Low",               # 最低价
        "成交量": "Volume",          # 成交量
        "成交额": "Turnover",        # 成交额
        "振幅": "Amplitude",         # 振幅
        "涨跌额": "PriceChange",     # 涨跌额
        "换手率": "TurnoverRate",    # 换手率
        "涨跌幅": "PriceChangePercentage",  # 涨跌幅
    }
    
    # 转换列名为英文
    data = transcolname(data, column_mapping)
    
    # 删除涨跌幅列（可能存在数据泄露问题）
    data.drop(columns=["PriceChangePercentage"], inplace=True)
    
    # 进行日期时间特征转换
    data = trans_datetime(data)

    return data


# 执行特征处理
feature = processing_feature()

# 输出处理后的特征数据
outputdata("./temp/feature.csv", feature)
